from pylab import *

CODE_PATH = '/home/xiaomin/wxm/Code/KaggleCCS'
caffe_root = '/opt/caffe/'  # this file should be run from {caffe_root}/examples (otherwise change this line)
data_root = '/home/xiaomin/wxm/Data/KaggleCCS/'
import sys

sys.path.insert(0, caffe_root + 'python')
import caffe

caffe.set_device(0)
caffe.set_mode_gpu()

NET_PATH_TEST = CODE_PATH + '/' + 'prototxt/vgg/VGG_2014_16_deploy.prototxt'
NET_PATH_TRAIN = CODE_PATH + '/' + 'prototxt/vgg/VGG_2014_16_train_2.prototxt'
WEIGHTS_PATH = '/home/xiaomin/wxm/Data/CaffeModels/VGG_ILSVRC_16_layers.caffemodel'
SOLVER_PATH = CODE_PATH + '/' + 'prototxt/vgg/solver.prototxt'
SNAPSHOTS_PATH = '/home/xiaomin/wxm/Data/KaggleCCS/snapshots'
LOG_FILE = '/home/xiaomin/wxm/Data/KaggleCCS/logs/vgg/2.txt'


def run_solver(niter, solver, disp_interval=10):
    """Run solvers for niter iterations,
       returning the loss and accuracy recorded each iteration.
       `solvers` is a list of (name, solver) tuples."""
    blobs = ('loss', 'accuracy_top1')
    (name, s) = solver
    loss, acc = ({name: np.zeros(niter)}
                 for _ in blobs)
    loss_disps = []

    log_file = open(LOG_FILE, 'w')

    for it in range(niter):
        s.step(1)  # run a single SGD step in Caffe
        # print s.net.blobs
        loss[name][it], acc[name][it] = (s.net.blobs[b].data.copy()
                                         for b in blobs)
        if it % disp_interval == 0 or it + 1 == niter:
            loss_disp = ''.join('iter:%s; loss=%.4f; acc=%.4f%%' %
                                (str(it), loss[name][it], np.round(100 * acc[name][it])))
            loss_disps.append(loss_disp)
            log_file.write(loss_disp + '\n')
            print loss_disp
            # Save the learned weights from both nets.
    # weight_dir = SNAPSHOTS_PATH
    # weights = {}
    # # name, s = solver
    # # filename = 'snapshot.caffemodel' % name
    # weights[name] = os.path.join(weight_dir, filename)
    # s.net.save(weights[name])
    return loss, acc


solver = caffe.get_solver(SOLVER_PATH)
solver.net.copy_from(WEIGHTS_PATH)

niter = 3000  # number of iterations to train

s = ('pretrained', solver)
print 'Running solvers for %d iterations...' % niter
loss, acc = run_solver(niter, s)
print 'Done.'

train_loss = loss['pretrained']
train_acc = acc['pretrained']


# Delete solvers to save memory.
del solver
